package video

import (
	"context"
	"encoding/json"
	"fmt"
	"gitee.com/kaylee595/gkaylee/kerr"
	"github.com/corona10/goimagehash"
	"github.com/disintegration/imaging"
	ffmpeg_go "github.com/u2takey/ffmpeg-go"
	"image"
	"image/color"
	"math"
	"os"
	"strconv"
)

type Info struct {
	Streams []InfoStream `json:"streams"`
	Format  struct {
		Filename       string `json:"filename"`
		NbStreams      int    `json:"nb_streams"`
		NbPrograms     int    `json:"nb_programs"`
		FormatName     string `json:"format_name"`
		FormatLongName string `json:"format_long_name"`
		StartTime      string `json:"start_time"`
		Duration       string `json:"duration"`
		Size           string `json:"size"`
		BitRate        string `json:"bit_rate"`
		ProbeScore     int    `json:"probe_score"`
		Tags           struct {
			MajorBrand       string `json:"major_brand"`
			MinorVersion     string `json:"minor_version"`
			CompatibleBrands string `json:"compatible_brands"`
			Encoder          string `json:"encoder"`
		} `json:"tags"`
	} `json:"format"`
}

type InfoStream struct {
	Index              int    `json:"index"`
	CodecName          string `json:"codec_name"`
	CodecLongName      string `json:"codec_long_name"`
	Profile            string `json:"profile"`
	CodecType          string `json:"codec_type"`
	CodecTimeBase      string `json:"codec_time_base"`
	CodecTagString     string `json:"codec_tag_string"`
	CodecTag           string `json:"codec_tag"`
	Width              int    `json:"width,omitempty"`
	Height             int    `json:"height,omitempty"`
	CodedWidth         int    `json:"coded_width,omitempty"`
	CodedHeight        int    `json:"coded_height,omitempty"`
	ClosedCaptions     int    `json:"closed_captions,omitempty"`
	HasBFrames         int    `json:"has_b_frames,omitempty"`
	SampleAspectRatio  string `json:"sample_aspect_ratio,omitempty"`
	DisplayAspectRatio string `json:"display_aspect_ratio,omitempty"`
	PixFmt             string `json:"pix_fmt,omitempty"`
	Level              int    `json:"level,omitempty"`
	ColorRange         string `json:"color_range,omitempty"`
	ColorSpace         string `json:"color_space,omitempty"`
	ColorTransfer      string `json:"color_transfer,omitempty"`
	ColorPrimaries     string `json:"color_primaries,omitempty"`
	ChromaLocation     string `json:"chroma_location,omitempty"`
	Refs               int    `json:"refs,omitempty"`
	IsAvc              string `json:"is_avc,omitempty"`
	NalLengthSize      string `json:"nal_length_size,omitempty"`
	RFrameRate         string `json:"r_frame_rate"`
	AvgFrameRate       string `json:"avg_frame_rate"`
	TimeBase           string `json:"time_base"`
	StartPts           int    `json:"start_pts"`
	StartTime          string `json:"start_time"`
	DurationTs         int    `json:"duration_ts"`
	Duration           string `json:"duration"`
	BitRate            string `json:"bit_rate"`
	BitsPerRawSample   string `json:"bits_per_raw_sample,omitempty"`
	NbFrames           string `json:"nb_frames"`
	Disposition        struct {
		Default         int `json:"default"`
		Dub             int `json:"dub"`
		Original        int `json:"original"`
		Comment         int `json:"comment"`
		Lyrics          int `json:"lyrics"`
		Karaoke         int `json:"karaoke"`
		Forced          int `json:"forced"`
		HearingImpaired int `json:"hearing_impaired"`
		VisualImpaired  int `json:"visual_impaired"`
		CleanEffects    int `json:"clean_effects"`
		AttachedPic     int `json:"attached_pic"`
		TimedThumbnails int `json:"timed_thumbnails"`
	} `json:"disposition"`
	Tags struct {
		Language    string `json:"language"`
		HandlerName string `json:"handler_name"`
		VendorId    string `json:"vendor_id"`
	} `json:"tags"`
	SampleFmt     string `json:"sample_fmt,omitempty"`
	SampleRate    string `json:"sample_rate,omitempty"`
	Channels      int    `json:"channels,omitempty"`
	ChannelLayout string `json:"channel_layout,omitempty"`
	BitsPerSample int    `json:"bits_per_sample,omitempty"`
	MaxBitRate    string `json:"max_bit_rate,omitempty"`
}

func Probe(filename string) (*Info, error) {
	s, err := ffmpeg_go.Probe(filename)
	if err != nil {
		return nil, kerr.WrapTraceStack(err)
	}
	i := new(Info)
	err = json.Unmarshal([]byte(s), &i)
	if err != nil {
		return nil, kerr.WrapTraceStack(err)
	}
	return i, nil
}

func combineImages(image1 image.Image, image2 image.Image) image.Image {
	i1Size := image1.Bounds().Size()
	i2Size := image2.Bounds().Size()
	canvasWidth := i1Size.X + i2Size.X
	canvasHeight := int(math.Max(float64(i1Size.Y), float64(i2Size.Y)))
	montage := imaging.New(canvasWidth, canvasHeight, color.NRGBA{})
	montage = imaging.Paste(montage, image1, image.Pt(0, 0))
	montage = imaging.Paste(montage, image2, image.Pt(i1Size.X, 0))
	return montage
}

func generateSprite(videoFilename string) (image.Image, error) {
	info, err := Probe(videoFilename)
	if err != nil {
		return nil, kerr.WrapTraceStack(err)
	}
	duration, err := strconv.ParseFloat(info.Format.Duration, 10)
	if err != nil {
		return nil, kerr.WrapTraceStack(err)
	}
	// ffmpeg -i videoFilename -ss 980 -vf scale=-1:'min(480,ih)' -vframes 1 -y previewImgFilename
	err = ffmpeg_go.Input(videoFilename, ffmpeg_go.KwArgs{"ss": int(duration * 0.05)}).Output("temp-p1.png", ffmpeg_go.KwArgs{
		"vf":      "scale='min(240,ih)':-1",
		"vframes": 1,
	}).Silent(true).OverWriteOutput().Run()
	if err != nil {
		return nil, kerr.WrapTraceStack(err)
	}
	defer os.Remove("temp-p1.png")
	err = ffmpeg_go.Input(videoFilename, ffmpeg_go.KwArgs{"ss": int(duration * 0.95)}).Output("temp-p2.png", ffmpeg_go.KwArgs{
		"vf":      "scale='min(240,ih)':-1",
		"vframes": 1,
	}).Silent(true).OverWriteOutput().Run()
	if err != nil {
		return nil, kerr.WrapTraceStack(err)
	}
	defer os.Remove("temp-p2.png")
	image1, err := imaging.Open("temp-p1.png")
	if err != nil {
		return nil, kerr.WrapTraceStack(err)
	}
	image2, err := imaging.Open("temp-p2.png")
	if err != nil {
		return nil, kerr.WrapTraceStack(err)
	}
	return combineImages(image1, image2), nil
}

func PHash(videoFilename string) (*goimagehash.ImageHash, error) {
	sprite, err := generateSprite(videoFilename)
	if err != nil {
		return nil, kerr.WrapTraceStack(err)
	}
	hash, err := goimagehash.PerceptionHash(sprite)
	if err != nil {
		return nil, kerr.WrapTraceStack(err)
	}
	return hash, nil
}

func GeneratePreviewVideo(ctx context.Context, input, output string) error {
	info, err := Probe(input)
	if err != nil {
		return kerr.WrapTraceStack(err)
	}
	duration, err := strconv.ParseFloat(info.Format.Duration, 10)
	if err != nil {
		return kerr.WrapTraceStack(err)
	}
	offset := float64(7)
	for i := 0; i < 10; i++ {
		sec := int(duration * (offset / 100))
		stream := ffmpeg_go.Input(input, ffmpeg_go.KwArgs{
			"ss": sec,
			"to": sec + 1,
		}).Silent(true)
		stream.Context = ctx
		err = stream.Output(fmt.Sprintf("temp-v%d.mp4", i), ffmpeg_go.KwArgs{"c:v": "copy", "an": ""}).OverWriteOutput().Run()
		if err != nil {
			return kerr.WrapTraceStack(err)
		}
		defer os.Remove(fmt.Sprintf("temp-v%d.mp4", i))
		offset += 10
	}
	streams := make([]*ffmpeg_go.Stream, 0)
	for i := 0; i < 10; i++ {
		streams = append(streams, ffmpeg_go.Input(fmt.Sprintf("temp-v%d.mp4", i)))
	}
	stream := ffmpeg_go.Concat(streams)
	stream.Context = ctx
	err = stream.Filter("scale", ffmpeg_go.Args{"'min(720,iw)':-2"}).
		Silent(true).Output(output, ffmpeg_go.KwArgs{"vsync": 2}).OverWriteOutput().Run()
	if err != nil {
		return kerr.WrapTraceStack(err)
	}
	return nil
}

func GeneratePreviewImg(ctx context.Context, videoFilename, previewImgFilename string) error {
	info, err := Probe(videoFilename)
	if err != nil {
		return kerr.WrapTraceStack(err)
	}
	duration, err := strconv.ParseFloat(info.Format.Duration, 10)
	if err != nil {
		return kerr.WrapTraceStack(err)
	}
	// ffmpeg -i videoFilename -ss 980 -vf scale='min(720,iw)':-2 -vframes 1 -y previewImgFilename
	// 获取视频 96% 位置的图片
	input := ffmpeg_go.Input(videoFilename, ffmpeg_go.KwArgs{"ss": int(duration * 0.96)}).Silent(true)
	input.Context = ctx
	err = input.Output(previewImgFilename, ffmpeg_go.KwArgs{
		"vf":      "scale='min(720,iw)':-2",
		"vframes": 1,
		"quality": 100,
	}).OverWriteOutput().Run()
	return kerr.WrapTraceStack(err)
}
